from __future__ import division
import numpy as np
#from numpy import multiply as nm
from ShallowDonor import Shallow_Donor
from Bulk import Silicon
from Constant_Physics import Physics
from Tools import *
from MRTools import *
from MRTools import Lorentzian as lorentz
import matplotlib.pyplot as plt
import time
from scipy.integrate import odeint
import src_fortran.froutines as fr
from termcolor import colored

#To Do list
#    - B sweep, with good linewidth
#    - modulation
#    - several transitions
#    - fitting parameters to a B/f sweep
#    - export (.csv/.ps/...)

class Parametrized_SDR_model():
    
    def __init__(self, bulk, specie, photoexcitation=1, mwW=1, cp=1/2,cap=1/2,rp=0.01, rap=1,B0=0.1, xyz=np.array([0, 0, 1]),
                  mod_freq=100000, mod_phase=0,mod_form="sin",mod_amp=0.001, temperature=10):
        """
        specie: define tje shallow donor specie ("P", "As", "Sb5/2", "Sb7/2", "Bi")
        D: Dissociation rate
        G: Generation rate
        C: Capture rate
        Rp: Recombination rate of parallel pairs
        Rap: Recombination rate of antiparallel pairs
        W: carrier generation due to illumination
        temperature: in Kelvin
        average_dist: average separation between one donor and one PRC, in 'nm'
        """
        self.CT=Computation_time(True)
        self.CT.TIC("Initialization")
        self.fortran_compiled_functions=True
        
        self.PE = photoexcitation
        self.Cp = cp
        self.Cap = cap 
        self.Rp = rp
        self.Rap = rap
        self.mw_w=mwW #the microwave intensity
        
        self.field=Bfield(B0=B0, xyz=xyz, mod_freq=mod_freq, mod_phase=mod_phase,mod_form=mod_form,mod_amp=mod_amp)
        
        self.width_MHz=5
        self.width_mT=1
        
        self.temperature = temperature
        self.steadystate_margin=1e-6
        self.is_steady_state_reached_flag=True
        self.is_steady_state_reached_print=False
        self.non_resonant_matrix_pop_conservative=True
        self.resonant_matrix_pop_conservative=True
        self.coupling_PRC_donor=True
        
        self.print_population=False

        self.donor = specie
        self.bulk = bulk
        self.PRCdim = 2
        self.PRCspin = 1 / 2
        self.electronicdim = 2
        self.nucleardim = int(2 * self.donor.nspin + 1)
        #-----------------
        self.dim_een = int(self.PRCdim * self.electronicdim * self.nucleardim) #number of states with the donor ionized
        self.dim_n = int(self.nucleardim) #number of states with both ionized
        self.full_dim = self.dim_een + self.dim_een + self.dim_n

        self.Mbj = np.zeros([self.dim_een, self.dim_een])
        self.Mjb = np.zeros([self.dim_een, self.dim_een])
        self.Mbj_el_wise_square = np.zeros([self.dim_een, self.dim_een])
        self.Mjb_el_wise_square = np.zeros([self.dim_een, self.dim_een])
        self.Mbi = np.zeros([self.dim_een, self.dim_een])
        self.Mib = np.zeros([self.dim_een, self.dim_een])
        #self.Mij = np.zeros([self.dim_een, self.dim_een])
        #self.Mji = np.zeros([self.dim_een, self.dim_een])
        #self.Mjk = np.zeros([self.dim_n, self.dim_een])
        #self.Mkj = np.zeros([self.dim_een, self.dim_n])
        #-----------------
        
        # Rate matrix: column names
            #[0...c1_e, c2_i=c1_e...c2_e, c3_i=c2_e...c3_e]
            #[                                                ]
            #[                                                ]
            #[                                                ]
            #[                                                ]
        self.c1_i=0
        self.c1_e=self.dim_een
        self.c2_i=self.c1_e
        self.c2_e=self.dim_een + self.dim_een
        self.c3_i=self.c2_e
        self.c3_e=self.full_dim+1
        
        self.timescale_without_resonance=1e-3
        self.time_intervale_without_resonance=1e-3*self.timescale_without_resonance
        self.timescale_with_resonance=1e-4
        self.time_intervale_with_resonance=1e-3*self.timescale_with_resonance
        self.conductivity=[]
        self.steadystate=[]
        self.resoningstate=[]
        self.steady_conductivity=0
        
        self.dipolar_transitions=[]    # [[B0,li-lj,EPR,SDR],...]


        self.norm=0
        
        #This is a marker to remember the time during the simulation
        self.last_time_step=0


        # Non Resonant rate equations
        self.NonResonantRates = np.zeros([self.full_dim, self.full_dim])
        # Resonant rate equations
        self.ResonantRates = np.zeros([self.full_dim, self.full_dim])
        # Non resonant b term
        self.population_away_from_resonance = np.zeros(self.dim_een)
        # Resonant b term
        self.population_at_resonance = np.zeros(self.dim_een)


        #Construction of the eigenvectors lists and hamiltonians matrix
        self.eigvalDpRz = np.zeros([self.dim_een])
        self.eigvalDzRz = np.zeros([self.dim_een])
        self.hamiltonianDzRz = np.zeros([self.dim_een, self.dim_een])
        self.eigvalDpRm = np.zeros([self.dim_n])
        self.CT.TOC("Initialization")
        
    def Non_Resonant_Rates_Evaluation(self):
        """
        Description:
            Calculate the Non Resonant rates matrix
            20120911,9am: execution time around 65 ms
        To do:
            -include the relaxation (and its dependence on the magnetic field)
        Input:
            - t: the time
        Output:
            -None
        Modification of properties:
            -self.NonResonantRates
        Dependences:
            -Properties:
                -self.full_dim
                -self.Cp
                -self.Cap
                -self.PE: photoexcitqtion rate
                -self.Rp
                -self.Rap
                -self.dim_een
                -self.dim_n
            -Class:
                -self.donor
                -self.bulk
        """
        
        self.CT.TIC("Non_Resonant_Rates_Evaluation")
        if self.fortran_compiled_functions:
            o=np.zeros([self.dim_een+self.dim_een+self.dim_n,self.dim_een+self.dim_een+self.dim_n])
            self.NonResonantRates=fr.nonresmat(o,
                                           self.Mbj_el_wise_square,
                                           self.dim_een+self.dim_een+self.dim_n,
                                           self.dim_een,
                                           self.bulk.dim_PRC,
                                           self.donor.dim_e,
                                           self.dim_n,
                                           self.Cp,
                                           self.Cap,
                                           self.Rp,
                                           self.Rap,
                                           self.PE)
        else:
            #Construct self.NRrates
                # Ci: assuming there is no coupling at this step (even if we write the transformation matrix)
            Calpha = (self.Cp - self.Cap) / 2
            Cbeta = (self.Cp + self.Cap) / 2
            Ci = np.zeros([self.dim_een, self.dim_een])
            Ci = Calpha * np.kron(np.kron(self.bulk.Sz, self.donor.Sz), self.donor.Id) + Cbeta * np.kron(np.kron(self.bulk.Sd, self.donor.Sd), self.donor.Id)
            
                # Cj
            Cj = np.zeros([self.dim_een, self.dim_een])
            v=np.kron(np.array([self.Cp,self.Cap,self.Cap,self.Cp]),np.ones([self.dim_n]))
            for k in range(self.dim_een):
                Cj[k,:]=(v*self.Mbj_el_wise_square[:,k]).T
                #print v.size
                #print (v*self.Mbj_el_wise_square[:,k]).size
                #print Cj[k,:].size
    
                # Rj
            Rj = np.zeros([self.dim_een, self.dim_een])
            R2=np.array([self.Rp,self.Rap,self.Rap,self.Rp])
            R3=np.ones([self.dim_n])
            R1=np.kron(R2,R3)
            for j in range(self.dim_een):
                Rj[j,j]=np.dot(R1.T,self.Mbj_el_wise_square[:,j])
    
                # Rk
            Rk = np.zeros([self.dim_n, self.dim_een])
            for j in range(self.dim_n):
                R4=np.zeros([self.dim_n])
                R4[j]=1
                R4=np.kron(R2,R4)
                for i in range(self.dim_een):
                    Rk[j,i]=np.dot(R4,self.Mbj_el_wise_square[:,i])
                
            
            
                # W: we assume there is no polarization of the excitation light ...
            PEi = self.PE/self.dim_een*np.kron(np.ones([self.PRCdim*self.electronicdim,1]),self.donor.Id)
            PEk = self.PE / self.dim_n * self.donor.Id
    
                # Construct the rates matrix
            self.Ci=Ci
            self.Cj=Cj
            self.Rj=Rj
            self.Rk=Rk
            self.PEk=PEk
            self.PEi=PEi
            self.NonResonantRates[self.c1_i:self.c1_e, self.c1_i:self.c1_e] = -Ci
            self.NonResonantRates[self.c2_i:self.c2_e, self.c1_i:self.c1_e] = Cj
            self.NonResonantRates[self.c2_i:self.c2_e, self.c2_i:self.c2_e] = -Rj
            self.NonResonantRates[self.c3_i:self.c3_e, self.c2_i:self.c2_e] = Rk
            self.NonResonantRates[self.c3_i:self.c3_e, self.c3_i:self.c3_e] = -PEk
            self.NonResonantRates[self.c1_i:self.c1_e, self.c3_i:self.c3_e] = PEi
            
            self.non_resonant_matrix_pop_conservative=Check_population_conservation(self.NonResonantRates,"Non Resonant Rates",self.print_population) and self.non_resonant_matrix_pop_conservative
        self.CT.TOC("Non_Resonant_Rates_Evaluation")

    def Resonant_Rates_Evaluation(self,excitation_freq=[[9076,5]]):
        """
        Description:
            Apply to the 'base_matrix' a transition between two levels 'T' (=[M1,b1,M2,b2,B(Tesla)]) at a rate 'w'
            20120911,9am: execution time around 100 ms
        Todo:
            -Upgrade for several irradiated frequencies at the same time!!!
        Input:
            -k: the transitions index in self.nonzero_i_index,self.nonzero_j_index
        Output:
            -None
        Modification of properties:
            -self.ResonantRates
        Dependence:
            -Properties:
                -self.mw_w
                -self.NonResonantRates
                -self.print_population
                -self.resonant_matrix_pop_conservative
            -Methods:
                -None
        """
        self.CT.TIC("Resonant_Rates_Evaluation")
        if self.fortran_compiled_functions:
            o=np.zeros([self.dim_een+self.dim_een+self.dim_n,self.dim_een+self.dim_een+self.dim_n])
            o=fr.nonresmat(o,self.Mbj_el_wise_square,
                           self.dim_een+self.dim_een+self.dim_n,
                           self.dim_een,
                           self.bulk.dim_PRC,
                           self.donor.dim_e,
                           self.dim_n,
                           self.Cp,
                           self.Cap,
                           self.Rp,
                           self.Rap,
                           self.PE)
            f=excitation_freq[0][0]
            w=excitation_freq[0][1]
            self.ResonantRates=fr.resmat(o,self.Dipolar_Transition_Matrix,
                                         f,w,
                                         self.mw_w,
                                         self.eigvalDzRz,
                                         self.dim_een+self.dim_een+self.dim_n,
                                         self.dim_een)
        else:
            #Print information
            self.Non_Resonant_Rates_Evaluation()
            self.ResonantRates=np.copy(self.NonResonantRates)
            for [freq,width] in excitation_freq:
                lorentzian = lorentz(derivative=0)
    #            for k in range(self.n_of_transitions):
    #                i_H=self.nonzero_i_index[k]
    #                j_H=self.nonzero_j_index[k]
    #                i_NRR=self.dim_een+i_H
    #                j_NRR=self.dim_een+j_H
    #                center = np.abs(self.eigvalDzRz[i_H]-self.eigvalDzRz[j_H])
    #                intensity=(self.Dipolar_Transition_Matrix[i_H,j_H])**2
    #                coeffs = intensity, center, width
                    
                for i in range(self.Dipolar_Transition_Matrix[:, 0].size):
                    for j in range(self.Dipolar_Transition_Matrix[0, :].size):
                        center = np.abs(self.eigvalDzRz[i] - self.eigvalDzRz[j])
                        coeffs = self.Dipolar_Transition_Matrix[i, j] ** 2, center, width
                        #coeffs = 1, center, width
                        evalpts = freq
                        dipole_transition_strength = lorentzian.Eval(coeffs, evalpts)
                        #print center,self.Dipolar_Transition_Matrix[i, j] ** 2
                    
                        i_NRR=self.dim_een+i
                        j_NRR=self.dim_een+j
                    
                        self.ResonantRates[i_NRR,i_NRR]+=-self.mw_w*dipole_transition_strength
                        self.ResonantRates[j_NRR,j_NRR]+=-self.mw_w*dipole_transition_strength
                        self.ResonantRates[i_NRR,j_NRR]+=self.mw_w*dipole_transition_strength
                        self.ResonantRates[j_NRR,i_NRR]+=self.mw_w*dipole_transition_strength
                        
                    #if dipole_transition_strength>1e-6:
                    #    print dipole_transition_strength
                    
            self.resonant_matrix_pop_conservative=Check_population_conservation(self.ResonantRates,"Resonant Rates",self.print_population) and self.resonant_matrix_pop_conservative
        self.CT.TOC("Resonant_Rates_Evaluation")
                
    def Temporal_evolution_without_resonance(self,starting_time_step=0):
        """
        Description:
            Resolve the time dependent problem without resonance (solving dN/dt=self.NonResonantRates*N)
        Todo:
            -None
        Input:
            -None
        Output:
            -None
        Modification of properties:
            -self.NonResonantRates
            -self.steady_conductivity
        Dependence:
            -Properties:
                -self.full_dim
                -self.c3_i,self.c3_e
                -self.dim_n
                -self.timescale_without_resonance,self.time_intervale_without_resonance
                -self.steadystate
                -self.conductivity
            -Methods:
                -self.Non_Resonant_Rates_Evaluation()
                -self.Compute_conductivity(y,t.size)
                -self.is_steady_state_reached()
        """
        self.CT.TIC("Temporal_evolution_without_resonance")
        #Print information
        print colored("Temporal evolution without any resonance:",'blue')
        
        self.previous_step_t=starting_time_step
        self.Do_Eig(time_step=starting_time_step)
        
        """
        print self.dim_een+self.dim_een+self.dim_n
        print self.dim_een
        print self.bulk.dim_PRC
        print self.donor.dim_e
        print self.dim_n
        print self.Cp
        print self.Cap
        print self.Rp
        print self.Rap
        print self.PE
        """
        self.Non_Resonant_Rates_Evaluation() #Compute the rates matrix for the non resonant case
        
        # preparing the initial state, no nuclear polarization
        
        y0=np.zeros([self.full_dim])
        y0[self.c3_i:self.c3_e]=1/self.dim_n*np.ones([self.dim_n])
        
        # defining the time scale
        self.t1=Arange(self.last_time_step,self.last_time_step+self.timescale_without_resonance,self.time_intervale_without_resonance)
        
        # defining the function dy/dt = f(y,t)
        def func(y,t0):
            #Don't know why, but some iterations don't increase the time step
            if self.previous_step_t<t0:
                self.Do_Eig(time_step=t0)
                self.Non_Resonant_Rates_Evaluation() #Compute the rates matrix for the non resonant case

            self.previous_step_t=t0
            Print_progress(t0,self.t1[-1],start=self.t1[0])
        
            return np.dot(self.NonResonantRates,y)

        
        # performing the ODEINT
        y=odeint(func,y0,self.t1,hmax=self.time_intervale_without_resonance)
        self.last_time_step=self.t1[-1]
        
        self.conductivity_0_ss=self.Compute_conductivity(y,self.t1.size)
        self.Check_list()
  
        #taking the last solution step (last line) and save it as steadystate, and storing the corresponding conductivity
        self.steadystate = y[-1,:]
        self.steady_conductivity=self.conductivity[-1]
        self.CT.TOC("Temporal_evolution_without_resonance")

    def Temporal_evolution_with_resonance(self,starting_time_step=0,excitation_freq=[[9076,5]]):
        """
        Description:
            Resolve the time dependent problem at resonance (solving dN/dt=self.ResonantRates*N)
        Todo:
            -None
        Input:
            -None
        Output:
            -None
        Modification of properties:
            -self.ResonantRates
            -self.resonant_conductivity
        Dependence:
            -Properties:
                -self.timescale_with_resonance,self.time_intervale_with_resonance
                -self.steadystate
                -self.conductivity
            -Methods:
                -self.Resonant_Rates_Evaluation()
                -self.Compute_conductivity(y,t.size)
                -self.is_steady_state_reached()
        """
        self.CT.TIC("Temporal_evolution_with_resonance")
        #Print information
        print colored("\nTemporal evolution with resonance(s):",'blue')

        #Compute the rates matrix for the resonant case
        # if self.coupling_PRC_donor==False,, we compute it at each time step (resource consuming?)
        self.previous_step_t=starting_time_step
        
        self.Do_Eig(time_step=starting_time_step)
                
        self.Resonant_Rates_Evaluation(excitation_freq=excitation_freq)
        
        # preparing the initial state, no nuclear polarization
        y0=self.steadystate
        
        # defining the time scale
        self.t2=Arange(self.last_time_step,self.last_time_step+self.timescale_with_resonance,self.time_intervale_with_resonance)

        # defining the function dy/dt = f(y,t)
        def func(y,t0):
            #Don't know why, but some iterations don't increase the time step
            if self.previous_step_t<t0:
                self.Do_Eig(time_step=t0)
                #Compute the rates matrix for the resonant case
                # if self.coupling_PRC_donor==False,, we compute it at each time step (resource consuming?)
                self.Resonant_Rates_Evaluation(excitation_freq=excitation_freq)
            self.previous_step_t=t0
            Print_progress(t0,self.t2[-1],start=self.t2[0])
            return np.dot(self.ResonantRates,y)
            
        # performing the ODEINT
        y=odeint(func,y0,self.t2,hmax=self.time_intervale_with_resonance)
        self.last_time_step=self.t2[-1]
        
        self.conductivity_ss_mwss=self.Compute_conductivity(y,self.t2.size)#Only the last row conductivity
        self.Check_list()
        self.resoningstate = y[-1,:]
        self.resonant_conductivity=self.conductivity[-1]
        self.CT.TOC("Temporal_evolution_with_resonance")
  
    def EigDpRz(self,bfield=np.array([0,0,1])):
        """
        Description:
            Define the Hamiltonian in the DpRz state
        Todo:
            -None
        Input:
            -B: magnetic field
        Output:
            -self.eigvalDpRz: matrix of the Hamiltonian in the case D+ R0
        Modification of properties:
            -self.hamiltonianDpRz
            -self.Mbi
            -self.Mib
        Dependence:
            -Properties:
                -self.dim_een
                -self.B
            -Class:
                -self.donor
                -self.bulk
        Comments:
            -relaxation supposed to be faster than others processes
        """
        #nuclear Zeeman and PRC
        #we write only the eigenvalues as it is already diagonalized Assuming B=Bz (here, avoid a diagonalization)
        self.hamiltonianDpRz=np.kron(np.kron(self.bulk.PRC_Hamiltonian_Ze(bfield),self.donor.Sd),self.donor.Id)
        self.hamiltonianDpRz+=np.kron(np.kron(self.bulk.Sd,self.donor.D_Hamiltonian_Ze(bfield)),self.donor.Id)
        self.hamiltonianDpRz+=np.kron(np.kron(self.bulk.Sd,self.donor.Sd),self.donor.D_Hamiltonian_Zn(bfield))
        self.Mbi = np.diag(np.ones([self.dim_een]))
        self.Mib = np.diag(np.ones([self.dim_een]))

    def EigDzRz(self,bfield=np.array([0,0,1])):
        """
        Description:
            Define the Hamiltonian in the DzRz state
        Todo:
            -None
        Input:
            -None
        Output:
            -self.eigvalDpRz: matrix of the Hamiltonian in the case D+ R0
        Modification of properties:
            -self.hamiltonianDpRz
            -self.Mbj
            -self.Mjb
            -self.Mbj_el_wise_square
            -self.Mjb_el_wise_square
        Dependence:
            -Properties:
                -self.dim_een
                -self.B
                -self.coupling_PRC_donor: Allows B(t) or not
            -Class:
                -self.donor
                -self.bulk
        Comments:
            -relaxation supposed to be faster than others processes
        """
        o=self.donor.D_Hamiltonian_Dz(bfield)
        self.hamiltonianDzRz = np.kron(self.bulk.Sd,o)    #Donor term
        self.hamiltonianDzRz += np.kron(self.bulk.PRC_Hamiltonian_Ze(bfield),np.kron(self.donor.Sd, self.donor.Id)) #PRC term
        
        if self.coupling_PRC_donor==True:
            #H=self.donor.D_Hamiltonian_Dz(self.field.field_vec(time_step))
            H=self.hamiltonianDzRz
            self.eigvalDzRz, self.Mbj = np.linalg.eig(H)  #Diagonalization
            self.Mjb = np.linalg.inv(self.Mbj)
            
        else:
            H=self.donor.D_Hamiltonian_Dz(bfield)
            def delta(H,i):
                return (H[i, i] - H[i + 9, i + 9])**2 + 4*H[i, i + 9]**2
            def p11(H,i):
                return -H[i, i + 9]
            def p12(H,i):
                return -H[i, i + 9]
            def p21(H, i):
                return (H[i, i] - H[i + 9, i + 9] - np.sqrt(delta(H, i)))/2
            def p22(H, i):
                return (H[i, i] - H[i + 9, i + 9] + np.sqrt(delta(H, i)))/2
            def ip11(H, i):
                return (-H[i, i] + H[i + 9, i + 9] - np.sqrt(delta(H, i)))/(2*H[i, i + 9]*np.sqrt(delta(H, i)))
            def ip12(H, i):
                return -1/np.sqrt(delta(H, i))
            def ip21(H, i):
                return (H[i, i] - H[i + 9, i + 9] - np.sqrt(delta(H, i)))/(2*H[i, i + 9]*np.sqrt(delta(H, i)))
            def ip22(H, i):
                return 1/np.sqrt(delta(H, i))
            n=len(H)
            self.Mbj=np.zeros([n,n])
            self.Mjb=np.zeros([n,n])
            self.Mbj[0,0]=self.Mbj[-1,-1]=1
            self.Mjb[0,0]=self.Mjb[-1,-1]=1
            for i in range(1,int(self.donor.nspin*2)+1):
                P11=p11(H,i)
                P21=p21(H,i)
                nx1=np.sqrt(P11**2+P21**2)
                P12=p12(H,i)
                P22=p22(H,i)
                nx2=np.sqrt(P12**2+P22**2)
                inx1=inx2=1
                self.Mbj[i,i]=P11/inx1
                self.Mbj[i,i+9]=P12/inx2
                self.Mbj[i+9,i]=P21/inx1
                self.Mbj[i+9,i+9]=P22/inx2
                
                iP11=ip11(H,i)
                iP21=ip21(H,i)
                inx1=np.sqrt(iP11**2+iP21**2)
                iP12=ip12(H,i)
                iP22=ip22(H,i)
                inx2=np.sqrt(iP12**2+iP22**2)
                inx1=inx2=1
                self.Mjb[i,i]=iP11/inx1
                self.Mjb[i,i+9]=iP12/inx2
                self.Mjb[i+9,i]=iP21/inx1
                self.Mjb[i+9,i+9]=iP22/inx2
                
            self.Mbj=np.kron(self.bulk.Sd,self.Mbj)
            self.Mjb=np.kron(self.bulk.Sd,self.Mjb)
            self.eigvalDzRz=np.diag(np.dot(np.dot(self.Mjb,np.kron(self.bulk.Sd,H)),self.Mbj))
            #hinton(1e-3*np.dot(np.dot(self.Mjb,np.kron(self.bulk.Sd,H)),self.Mbj))
            #print self.eigvalDzRz
        self.Mbj_el_wise_square=self.Mbj**2
        self.Mjb_el_wise_square=self.Mjb**2

    def EigDpRm(self,bfield=np.array([0,0,1])):
        """
        Description:
            Define the Hamiltonian in the DzRz state
        Todo:
            -None
        Input:
            -B: magnetic field
        Output:
            -None
        Modification of properties:
            -self.hamiltonianDpRm
            -self.Mkb
            -self.Mbk
        Dependence:
            -Properties:
                -self.dim_een
                -self.B
            -Class:
                -self.donor
                -self.bulk
        Comments:
            -relaxation supposed to be faster than others processes
        """
        #nuclear Zeeman and PRC
        #we write only the eigenvalues as it is already diagonalized
        self.hamiltonianDpRm=np.kron(np.kron(self.bulk.Sd,self.donor.Sd),self.donor.D_Hamiltonian_Zn(bfield))
        
    def Do_Eig(self,time_step=0):
        """
        Description:
            Perform the eigensystem, fill the transitions list...
            20120911,9am: execution time around 2 ms
        Todo:
            -None
        Input:
            -B: from MRTools.Bfield(B,xyz)
        Output:
            -None
        Modification of properties:
            -self.Dipolar_Transition_Matrix
            -self.n_of_transitions
            -self.dipolar_transitions
            -self.nonzero_i_index
            -self.nonzero_j_index
        Dependence:
            -Properties:
                -self.Mjb
                -self.Mbj
                -self.Dipolar_Transition_Matrix
            -Methods:
                -self._EigDpRz()
                -self._EigDzRz()
                -self._EigDpRm()
            -class:
                -self.donor
                -self.bulk
        """
        self.CT.TIC("Do_Eig")
        bfield=self.field.field_vec(time_step)
        self.EigDpRz(bfield=bfield)
        self.EigDzRz(bfield=bfield)
        self.EigDpRm(bfield=bfield)
        #self.Mij = np.dot(self.Mib, self.Mbj)
        #self.Mji = np.dot(self.Mjb, self.Mbi)
        self.CT.TOC("Do_Eig")
        self.CT.TIC("Do_Eig2")
        DP1=np.kron(np.kron(self.bulk.Sd,self.donor.Sx),self.donor.Id)    #Dipolar transition
        self.Dipolar_Transition_Matrix=np.dot(self.Mjb,np.dot(DP1,self.Mbj))   #Change the basis
        self.n_of_transitions=np.count_nonzero(self.Dipolar_Transition_Matrix)
        self.dipolar_transitions=np.zeros([self.n_of_transitions,4])    # [[B0,li-lj,EPR],...]
        self.nonzero_i_index,self.nonzero_j_index=np.nonzero(self.Dipolar_Transition_Matrix)
        self.CT.TOC("Do_Eig2")

    def Compute_conductivity(self,odeint_sol,time_scale):
        """
        Description:
            Compute the conductivity knowing the populations
        Todo:
            -None
        Input:
            -odeint_sol: solution of the system
            -time_scale=t.size
        Output:
            -self.conductivity
        Modification of properties:
            -self.conductivity
        Dependence:
            -Properties:
                -self.c1_i,self.c1_e
                -self.c3_i,self.c3_e
            -class:
                -self.donor
                -self.bulk
        
        """
        self.CT.TIC("Compute_conductivity")
        self.conductivity=np.zeros([time_scale])
        for k in range(time_scale):
            self.conductivity[k]=self.bulk.mue*np.sum(odeint_sol[k,self.c1_i:self.c1_e])
            self.conductivity[k]+=self.bulk.muh*np.sum(odeint_sol[k,self.c3_i:self.c3_e])
        return self.conductivity
        self.CT.TOC("Compute_conductivity")
        
    def is_steady_state_reached(self):
        """
        Description:
            Print whether the steady state is reached or not
        Todo:
            -Redefine the steady state
        Input:
            -odeint_sol: solution of the system
            -time_scale=t.size
        Output:
            -None
        Modification of properties:
            -None
        Dependence:
            -Properties:
                -self.conductivity
                -self.steadystate_margin
                -self.is_steady_state_reached_print
        """
        self.CT.TIC("is_steady_state_reached")
        if self.conductivity==[]:
            print "\nPlease compute the conductivity first."
        else:
            if np.abs(self.conductivity[-1]-self.conductivity[-2])<=self.steadystate_margin:
                if self.is_steady_state_reached_print==True:
                    print "\nSteady state reached: ds/dt= ",np.abs(self.conductivity[-1]-self.conductivity[-2])
            else:
                if self.is_steady_state_reached_print==True:
                    print "\nSteady state not reached yet: ds/dt= ",np.abs(self.conductivity[-1]-self.conductivity[-2])
                self.is_steady_state_reached_flag=False
        self.CT.TOC("is_steady_state_reached")
   
    def Plot_transient(self,transient_type="conductivity,light and resonance transient"):
        """
        Input:
            - transient_type:    "conductivity": plot the conductivity transient
        """
        
        self.CT.TIC("Plot_transient")
        #Determing the number of plots
        if "light" in transient_type and "resonance" in transient_type:
            n_plots=2
        elif "light" in transient_type and "resonance" not in transient_type:
            n_plots=1
        
        if "light" in transient_type:
            fig = plt.figure()
            ax=fig.add_subplot(n_plots,1,1)
            ax.plot(self.t1,self.conductivity_0_ss)
            ax.set_title('Light on - transient')
            
        if "resonance" in transient_type:
            ax=fig.add_subplot(n_plots,1,2)
            ax.plot(self.t2,self.conductivity_ss_mwss)
            ax.set_title('Resonance on - transient')
            
        plt.show()
        self.CT.TOC("Plot_transient")
          
    def Check_list(self):
        """
        Description:
            Print whether the matrixes are population conservative or not, and if the steady state is reached or not
        Todo:
            -None
        Input:
            -None
        Output:
            -None
        Modification of properties:
            -None
        Dependence:
            -Properties:
                -self.non_resonant_matrix_pop_conservative
                -self.resonant_matrix_pop_conservative
                -self.is_steady_state_reached_flag
        """
        self.CT.TIC("Check_list")
        if not self.non_resonant_matrix_pop_conservative:
            print "\nThe non resonant matrix does not conserve the population."
        
        if not self.resonant_matrix_pop_conservative:
            print "\nAt least one of the resonant matrix does not conserve the population."
        
        if not self.is_steady_state_reached_flag:
            print "\nOne of the steady states at least was not reached."
        self.CT.TOC("Check_list")
  
    def Plot_populations(self,ploted_variable="steady state population",plot_size="full"):
        """
        Input:
            - ploted_variable=      "steady state populations"
                                    "resoning state populations"
                                    "difference populations"
            - plot_size =   "full"
                            "donor"
        """
        self.CT.TIC("Plot_populations")
        if plot_size=="full":
            ind1 = range(self.full_dim)  # the x locations for the groups
        elif plot_size=="donor":
            ind1 = range(self.c2_i,self.c2_e)  # the x locations for the groups
            
        w = 0.4       # the width of the bars
        y1=[]
        for k in ind1:
            if ploted_variable =="steady state populations":
                y1.append(self.steadystate[k])
                title="Steady state populations"
            elif ploted_variable =="resoning state populations":
                y1.append(self.resoningstate[k])
                title="Populations at resonance"
            elif ploted_variable =="difference populations":
                y1.append(self.steadystate[k]-self.resoningstate[k])
                title="Populations at SS - at resonance"
        
        fig=plt.figure()
        ax = fig.add_subplot(1,1,1)
        
        ax.bar(ind1, y1, width=w,color='r')
        
        # add some
        plt.ylabel('Populations')
        plt.title(title)
        
        plt.show()
        self.CT.TOC("Plot_populations")

    def F_sweep_EPR(self):
        """
        Checked
        """
        
        self.CT.TIC("F_sweep_EPR")
        width = 10 #MHz
        n_points = 10000
        Y = np.zeros(n_points)
        X = np.arange(0, n_points, 1)

        ham = self.donor.D_Hamiltonian_Dz(self.field.field_vec(0))
        eigvalDz, eigvecDz = np.linalg.eig(ham)  #Diagonalization
        P = eigvecDz
        iP = np.linalg.inv(P)
        MR = np.kron(self.donor.Sx, self.donor.Id)
        MR_Mb = np.dot(iP, np.dot(MR, P))

        lorentzian = lorentz()
        for i in range(MR_Mb[:, 0].size):
            for j in range(MR_Mb[0, :].size):
                center = np.abs(eigvalDz[i] - eigvalDz[j])
                coeffs = MR_Mb[i, j] ** 2, center, width
                evalpts = X
                Y += lorentzian.Eval(coeffs, evalpts)
                print center,MR_Mb[i, j] ** 2
                
        plt.plot(X, Y, '-')
        plt.show()
        self.CT.TOC("F_sweep_EPR")

if __name__ == "__main__":
    #Instanciation
    tic=time.clock()
    Bi = Shallow_Donor("Bi")
    P = Shallow_Donor("P")
    Si = Silicon()
    
    SDR = Parametrized_SDR_model(Si, Bi,photoexcitation=1e3,mwW=1e10, 
                                 rp=2e4, rap=1e6, #Hz
                                 cp=1e4, cap=1e4, #Hz
                                 B0=0.066056, xyz=np.array([0, 0, 1]), #T, 0.066056
                                 mod_freq=100000, mod_phase=0,mod_form="sin",mod_amp=0.001) #Hz, T
    SDR.timescale_without_resonance=1e-3
    SDR.timescale_with_resonance=1e-3
    SDR.Temporal_evolution_without_resonance()
    SDR.Temporal_evolution_with_resonance(excitation_freq=[[9076,1]]) #9076
    SDR.Plot_transient(transient_type="conductivity,light and resonance transient")# and resonance transient
    
    #SDR.F_sweep_EPR()
    #SDR.coupling_PRC_donor=False
    #SDR.EigDzRz(0)
    
    #SDR.Plot_populations(ploted_variable="difference populations",plot_size="full")#"donor","full"

    toc=time.clock()
    print "\nOperation finished. Processing time: ",toc-tic,"\n"
    SDR.CT.Display()